import pandas as pd
from typing import Union, Dict, Any, List, Optional
class DFDataEngine:
    def __init__(
        self,
        data: Union[Dict[str, pd.DataFrame], pd.DataFrame, List[Dict[str, Any]]],
        table_name: Optional[str] = None,
        return_format: str = "json",
        auto_create_table: bool = True,
        auto_add_id: bool = True,
    ):
        """
        初始化 DFDataEngine。
        :param data: 数据源，可以是字典、DataFrame 或列表。
        :param table_name: 默认表名。
        :param return_format: 返回格式（json 或 DataFrame）。
        :param auto_create_table: 是否自动创建表。
        :param auto_add_id: 是否自动添加 ID 列。
        """
        self.return_format = return_format
        self.auto_create_table = auto_create_table
        self.auto_add_id = auto_add_id
        self.id_column = "id"  # 默认主键列名
        self.default_table_name = table_name or "default"  # 默认表名

        # 初始化 tables
        self.tables = self._convert_to_tables(data)

    def _convert_to_tables(self, data: Union[Dict[str, pd.DataFrame], pd.DataFrame, List[Dict[str, Any]]]) -> Dict[str, pd.DataFrame]:
        """
        将输入数据转换为表格字典。
        """
        if isinstance(data, dict):
            # 如果 data 是字典，直接使用
            tables = data
            if self.auto_add_id:
                for table_name, df in tables.items():
                    if self.id_column not in df.columns:
                        df[self.id_column] = range(1, len(df) + 1)
            return tables
        elif isinstance(data, pd.DataFrame):
            # 如果 data 是 DataFrame，转换为字典
            df = data
            if self.auto_add_id and self.id_column not in df.columns:
                df[self.id_column] = range(1, len(df) + 1)
            return {self.default_table_name: df}
        elif isinstance(data, list):
            # 如果 data 是列表，转换为 DataFrame
            df = pd.DataFrame(data)
            if self.auto_add_id and self.id_column not in df.columns:
                df[self.id_column] = range(1, len(df) + 1)
            return {self.default_table_name: df}
        else:
            raise ValueError("不支持的数据类型。请提供字典、DataFrame 或列表。")

    def save(self, entity: Dict[str, Any], table_name: Optional[str] = None) -> str:
        """
        插入数据。
        """
        if table_name is None:
            table_name = self.default_table_name

        if table_name not in self.tables:
            if self.auto_create_table:
                columns = [self.id_column] if self.auto_add_id else []
                self.tables[table_name] = pd.DataFrame(columns=columns)
            else:
                raise ValueError(f"表 '{table_name}' 不存在，且自动建表已关闭。")

        df = self.tables[table_name]
        new_row = pd.DataFrame([entity])
        df = pd.concat([df, new_row], ignore_index=True)

        if self.auto_add_id and self.id_column in df.columns:
            if df[self.id_column].isnull().any():
                max_id = df[self.id_column].max(skipna=True)
                if pd.isna(max_id):
                    max_id = 0
                df.loc[df[self.id_column].isnull(), self.id_column] = range(int(max_id) + 1, int(max_id) + 1 + df[self.id_column].isnull().sum())

        self.tables[table_name] = df
        return str(df.iloc[-1][self.id_column])

    def delete(self, query: Dict[str, Any], table_name: Optional[str] = None) -> int:
        """
        删除数据。
        """
        if table_name is None:
            table_name = self.default_table_name

        if table_name not in self.tables:
            raise ValueError(f"表 '{table_name}' 不存在。")

        df = self.tables[table_name]
        for key, value in query.items():
            df = df[df[key] != value]
        deleted_count = len(self.tables[table_name]) - len(df)
        self.tables[table_name] = df
        return deleted_count

    def deleteById(self, id: str, table_name: Optional[str] = None) -> int:
        """
        根据 ID 删除数据。
        """
        if table_name is None:
            table_name = self.default_table_name

        if table_name not in self.tables:
            raise ValueError(f"表 '{table_name}' 不存在。")

        df = self.tables[table_name]
        if self.id_column not in df.columns:
            raise ValueError(f"表 '{table_name}' 中没有 '{self.id_column}' 列。")

        df = df[df[self.id_column] != id]
        deleted_count = len(self.tables[table_name]) - len(df)
        self.tables[table_name] = df
        return deleted_count

    def update(self, query: Dict[str, Any], update_data: Dict[str, Any], table_name: Optional[str] = None) -> int:
        """
        更新数据。
        """
        if table_name is None:
            table_name = self.default_table_name

        if table_name not in self.tables:
            raise ValueError(f"表 '{table_name}' 不存在。")

        df = self.tables[table_name]
        for key, value in query.items():
            df.loc[df[key] == value, list(update_data.keys())] = list(update_data.values())
        self.tables[table_name] = df
        return len(df[df[key] == value])

    def updateById(self, id: str, update_data: Dict[str, Any], table_name: Optional[str] = None) -> int:
        """
        根据 ID 更新数据。
        """
        if table_name is None:
            table_name = self.default_table_name

        if table_name not in self.tables:
            raise ValueError(f"表 '{table_name}' 不存在。")

        df = self.tables[table_name]
        if self.id_column not in df.columns:
            raise ValueError(f"表 '{table_name}' 中没有 '{self.id_column}' 列。")

        df.loc[df[self.id_column] == id, list(update_data.keys())] = list(update_data.values())
        self.tables[table_name] = df
        return 1

    def find(self, query: Dict[str, Any], table_name: Optional[str] = None) -> Union[List[Dict], pd.DataFrame]:
        """
        查询数据。
        :param query: 查询条件字典。
        :param table_name: 表名。
        :return: 查询结果（列表或 DataFrame）。
        :raises ValueError: 如果表不存在或字段不存在。
        """
        if table_name is None:
            table_name = self.default_table_name

        # 检查表是否存在
        if table_name not in self.tables:
            raise ValueError(f"表 '{table_name}' 不存在。")

        df = self.tables[table_name]

        # 检查查询条件中的字段是否存在
        for key in query.keys():
            if key not in df.columns:
                raise ValueError(f"表 '{table_name}' 中没有 '{key}' 列。")

        # 遍历查询条件，过滤数据
        for key, value in query.items():
            df = df[df[key] == value]

        if self.return_format == "json":
            return df.to_dict(orient="records")
        else:
            return df

    def findAll(self, table_name: Optional[str] = None) -> Union[List[Dict], pd.DataFrame]:
        """
        查询所有数据。
        """
        if table_name is None:
            table_name = self.default_table_name

        if table_name not in self.tables:
            raise ValueError(f"表 '{table_name}' 不存在。")

        df = self.tables[table_name]
        if self.return_format == "json":
            return df.to_dict(orient="records")
        else:
            return df

    def findById(self, id: str, table_name: Optional[str] = None) -> Union[Dict, pd.DataFrame]:
        """
        根据 ID 查询数据。
        """
        if table_name is None:
            table_name = self.default_table_name

        if table_name not in self.tables:
            raise ValueError(f"表 '{table_name}' 不存在。")

        df = self.tables[table_name]
        if self.id_column not in df.columns:
            raise ValueError(f"表 '{table_name}' 中没有 '{self.id_column}' 列。")

        result = df[df[self.id_column] == id]
        if result.empty:
            return None

        if self.return_format == "json":
            return result.iloc[0].to_dict()
        else:
            return result

    def findPage(self, query: Dict[str, Any], page: int, page_size: int, table_name: Optional[str] = None) -> Union[List[Dict], pd.DataFrame]:
        """
        分页查询数据。
        """
        if table_name is None:
            table_name = self.default_table_name

        if table_name not in self.tables:
            raise ValueError(f"表 '{table_name}' 不存在。")

        df = self.tables[table_name]
        for key, value in query.items():
            df = df[df[key] == value]

        skip = (page - 1) * page_size
        df = df.iloc[skip:skip + page_size]

        if self.return_format == "json":
            return df.to_dict(orient="records")
        else:
            return df


